package com.turing.draw.utils;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.turing.common.entity.JavaOcrParses;
import com.turing.common.entity.JavaTableParses;
import com.turing.common.entity.ParseContentVO;
import com.turing.common.entity.ParseTableVO;
import com.turing.common.util.FileUtils;
import com.turing.common.util.Log;
import com.turing.common.util.StringUtils;
import com.turing.draw.vo.in.OcrResultEntity;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.CollectionUtils;

import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

/**
 * @ClassName TableUtils
 * @Description TODO
 * @Author Distance
 * @Date 2022/11/21 10:18
 * @Version 1.0
 **/
public class TableUtil {
    private static final Logger logger = LoggerFactory.getLogger(TableUtil.class);

    /**
     * 中文正则
     */
    private static final Pattern MATCH_CHINESE = Pattern.compile("[\u4e00-\u9fa5]+");

    /**
     * ocr表格识别分隔符
     */
    private static final String OCR_TABLE_SPLIT = "++++++++++++++++";

    /**
     * 需要获取的表格往上的文本行数
     */
    private static final Integer TEXT_LINE_LIMIT = 10;

    /**
     * 需要进行表格合并的最小表格数
     */
    private static final Integer NEED_MERGE_MIN_NUM = 2;

    /**
     * 单元格列或行坐标的分割数目
     */
    private static final Integer CELL_SPILT_NUM = 2;

    public static List<ParseTableVO> handleTables(JSONArray json) {
        List<ParseTableVO> tableVOList = new ArrayList<>();
        if (json == null || json.isEmpty()) {
            logger.error("=========> 处理表格传进来的json为空 <=============");
            return tableVOList;
        }

        //获取ocr结构中的表格数据 map（key：表格id，value：表格数据）
        Map<String, List<OcrResultEntity.PageEntity.LineEntity>> tableMap = getTables(json);
        //对抽取出来的表格数据进行进一步处理
        tableVOList = doHandleTables(json, tableMap);
        return tableVOList;
    }

    /**
     * 从ocr json中提取表格数据
     *
     * @param json
     * @return
     */
    private static Map<String, List<OcrResultEntity.PageEntity.LineEntity>> getTables(JSONArray json) {
        Map<String, List<OcrResultEntity.PageEntity.LineEntity>> tableMap = new HashMap<>(5);


        //遍历ocr json每一页的数据
        for (Object pageJson : json) {
            OcrResultEntity.PageEntity pageEntity = JSONObject.parseObject(JSONObject.toJSONString(pageJson), OcrResultEntity.PageEntity.class);
            //页码
            Integer pageNo = pageEntity.getPageNo();
            //获取行数据
            List<OcrResultEntity.PageEntity.LineEntity> lineList = pageEntity.getLineList();
            //用于存储表格数据的map集合，key为表格id，value为单元格集合
            Map<Integer, List<OcrResultEntity.PageEntity.LineEntity>> map = new HashMap<>(5);

            //遍历循环每一行数据，提取单元格
            for (OcrResultEntity.PageEntity.LineEntity lineEntity : lineList) {
                List<OcrResultEntity.CellInfo> cells = lineEntity.getCells();
                if (cells == null || cells.isEmpty()) {
                    continue;
                }

                //表格id
                Integer id = lineEntity.getTableId();
                //获取map中改表格id的line列表，没有则创建一个新的list并返回
                List<OcrResultEntity.PageEntity.LineEntity> list = map.computeIfAbsent(id, k -> new ArrayList<>());

                int index = 0;
                for (OcrResultEntity.CellInfo cell : cells) {
                    if (cell == null) {
                        continue;
                    }
                    //新建行对象，将单元格封装为行
                    OcrResultEntity.PageEntity.LineEntity tableLine = new OcrResultEntity.PageEntity.LineEntity();
                    tableLine.setLineId(lineEntity.getLineId());
                    tableLine.setLineNo(lineEntity.getLineNo());
                    tableLine.setObjContent(cell.getObjContent());
                    tableLine.setObjType(lineEntity.getObjType());
                    tableLine.setRowNo(lineEntity.getRowNo());
                    tableLine.setSortNo(lineEntity.getSortNo());
                    tableLine.setTableId(id);
                    //单元格起始行
                    tableLine.setRow_start(cell.getRow_start_end().get(0));
                    //单元格结束行
                    tableLine.setRow_end(cell.getRow_start_end().get(1));
                    //单元格起始列
                    tableLine.setCol_start(cell.getCol_start_end().get(0));
                    //单元格结束列
                    tableLine.setCol_end(cell.getCol_start_end().get(1));

                    //计算当前文本在这一行的文本位置
                    //文本开始索引
                    int start = lineEntity.getObjContent().indexOf(cell.getObjContent(), index);
                    //文本结束索引
                    int end = start + cell.getObjContent().length();
                    tableLine.setLinePos(Arrays.asList(start, end));
                    //设置indexOf索引的起始位置
                    index = end;

                    //设置坐标信息
                    JSONArray objPos = new JSONArray();
                    objPos.addAll(cell.getObjPos());
                    tableLine.setObjPos(objPos);
                    list.add(tableLine);
                }
            }

            if (!map.isEmpty()) {
                //重新设置tableId，因为ocr是按图片识别，每一页识别一次，所以tableId都是以0开头
                List<Integer> ids = map.keySet().stream().sorted().collect(Collectors.toList());
                //为封装出来的单元格行数据重设tableId
                for (Integer id : ids) {
                    String finTableId = pageNo + StringUtils.UNDERLINE + id;
                    List<OcrResultEntity.PageEntity.LineEntity> lines = map.get(id);
                    for (OcrResultEntity.PageEntity.LineEntity line : lines) {
                        line.setFinTableId(finTableId);
                    }
                    tableMap.put(finTableId, lines);
                }
            }
        }

        return tableMap;
    }

    /**
     * 处理表格数据，封装成结构化数据
     *
     * @param json
     * @param tableMap
     * @return
     */
    private static List<ParseTableVO> doHandleTables(JSONArray json, Map<String, List<OcrResultEntity.PageEntity.LineEntity>> tableMap) {
        //每一页json数据
        List<OcrResultEntity.PageEntity> pageEntities = json.toJavaList(OcrResultEntity.PageEntity.class);
        //获取所有页的行数据的lineId
        List<String> lineIds = pageEntities.stream().map(OcrResultEntity.PageEntity::getLineList).flatMap(Collection::stream).map(OcrResultEntity.PageEntity.LineEntity::getLineId).collect(Collectors.toList());
        //根据lineId对行进行分组
        Map<String, OcrResultEntity.PageEntity.LineEntity> lineMap = pageEntities.stream().map(OcrResultEntity.PageEntity::getLineList).flatMap(Collection::stream).collect(Collectors.toMap(OcrResultEntity.PageEntity.LineEntity::getLineId, line -> line));
        //tableId集合
        List<String> tableIds = tableMap.keySet().stream().sorted(TableUtil::compareFinTableId).collect(Collectors.toList());

        List<JavaTableParses> parseTableList = new ArrayList<>();

        for (String finTableId : tableIds) {
            JavaTableParses javaTableParses = new JavaTableParses();
            List<OcrResultEntity.PageEntity.LineEntity> table = tableMap.get(finTableId);
            List<JavaTableParses.JavaTableParse> javaTableParseList = new ArrayList<>();
            //封装成后处理需要的对象
            for (OcrResultEntity.PageEntity.LineEntity cell : table) {
                JavaTableParses.JavaTableParse javaTableParse = new JavaTableParses.JavaTableParse();
                javaTableParse.setRow(cell.getRow_start() + StringUtils.HORIZONTAL_BAR + (cell.getRow_end() - 1));
                javaTableParse.setColumn(cell.getCol_start() + StringUtils.HORIZONTAL_BAR + (cell.getCol_end() - 1));
                javaTableParse.setValue(cell.getObjContent());
                javaTableParse.setLineId(cell.getLineId());
                javaTableParse.setLinePos(cell.getLinePos());
                javaTableParseList.add(javaTableParse);
            }

            //获取表格前10行的文本信息
            List<JavaOcrParses> javaOcrParsesList = new ArrayList<>();
            String lineId = table.get(0).getLineId();
            int index = lineIds.indexOf(lineId);
            for (int i = index - 1; i >= Math.max(0, index - TEXT_LINE_LIMIT); i--) {
                String id = lineIds.get(i);
                OcrResultEntity.PageEntity.LineEntity lineEntity = lineMap.get(id);
                //遇到表格就终止
                if ("table".equals(lineEntity.getObjType()) || "table_postpreprocess".equals(lineEntity.getObjType_postpreprocess())) {
                    break;
                }
                //封装文本对象
                JavaOcrParses javaOcrParses = new JavaOcrParses();
                javaOcrParses.setLineId(id);
                javaOcrParses.setObjType(lineEntity.getObjType());
                javaOcrParses.setObjContent(lineEntity.getObjContent());
                javaOcrParses.setLineNo(lineEntity.getLineNo());
                JSONArray array = new JSONArray();
                array.addAll(lineEntity.getObjPos());
                javaOcrParses.setObjPos(array);
                javaOcrParsesList.add(javaOcrParses);
            }
            //对收集到的文本对象进行反转处理
            Collections.reverse(javaOcrParsesList);
            javaTableParses.setTable(javaTableParseList);
            javaTableParses.setOcr(javaOcrParsesList);

            parseTableList.add(javaTableParses);
        }

        //处理表格，表格合并等
        return parseTable(parseTableList, json.toJSONString());
    }

    /**
     * 比较finTableId大小
     *
     * @param id1
     * @param id2
     * @return
     */
    private static int compareFinTableId(String id1, String id2) {
        String[] split1 = id1.split(StringUtils.UNDERLINE);
        String[] split2 = id2.split(StringUtils.UNDERLINE);
        int page1 = Integer.parseInt(split1[0]);
        int page2 = Integer.parseInt(split2[0]);
        int tableId1 = Integer.parseInt(split1[1]);
        int tableId2 = Integer.parseInt(split2[1]);

        return Integer.compare(page1 + tableId1, page2 + tableId2);
    }

    /**
     * 合并表格
     *
     * @param ocrJson   ocr数据
     * @param tableList 表格列表
     * @return 合并之后的表格列表
     */
    public static List<JavaTableParses> mergeTableData(String ocrJson, List<JavaTableParses> tableList) {
        // 如果表格列表为空或者数量小于2，不需要合并表格
        if (CollectionUtils.isEmpty(tableList) || StringUtils.isEmpty(ocrJson)) {
            return new ArrayList<>();
        }

        if(tableList.size() < NEED_MERGE_MIN_NUM) {
            return tableList;
        }

        // 获取ocr数据中的所有行数据列表
        // OCR中一个表格结束，都会给一个++++++++++++, 因为模型的原因，有时表格结束会有，有的会没有，所以不考虑这行数据
        List<com.turing.common.entity.OcrResultEntity.PageEntity> ocrResultVOList = JSONArray.parseArray(ocrJson, com.turing.common.entity.OcrResultEntity.PageEntity.class);
        List<com.turing.common.entity.OcrResultEntity.PageEntity.LineEntity> tableLineList = ocrResultVOList.stream().map(com.turing.common.entity.OcrResultEntity.PageEntity::getLineList)
                .flatMap(Collection::stream).filter(e -> !e.getObjContent().equals(OCR_TABLE_SPLIT)).collect(Collectors.toList());

        // 记录每个行数据所在的索引
        Map<String, Integer> lineIndexMap = new HashMap<>(tableLineList.size());
        for (int i = 0; i < tableLineList.size(); i++) {
            lineIndexMap.put(tableLineList.get(i).getLineId(), i);
        }

        // 已经合并的表格列表, 先加入第一个表格
        List<JavaTableParses> mergedTableList = new ArrayList<>();
        mergedTableList.add(tableList.get(0));

        boolean needMerge;

        for (int i = 1; i < tableList.size(); i++) {
            // 已合并表格列表中的最后一个表格
            List<JavaTableParses.JavaTableParse> lastTable = mergedTableList.get(mergedTableList.size() - 1).getTable();
            // 当前未合并的表格
            List<JavaTableParses.JavaTableParse> curTable = tableList.get(i).getTable();

            // 最后一个表格的最后一个单元格
            JavaTableParses.JavaTableParse lastJavaTableParse = lastTable.get(lastTable.size() - 1);
            // 当前未合并的表格的第一个单元格
            JavaTableParses.JavaTableParse firstJavaTableParse = curTable.get(0);
            // 获取这2个单元格所在的页码
            int page1 = Integer.parseInt(lastJavaTableParse.getLineId().split(com.turing.common.util.StringUtils.UNDERLINE)[1]);
            int page2 = Integer.parseInt(firstJavaTableParse.getLineId().split(com.turing.common.util.StringUtils.UNDERLINE)[1]);
            // 获取第一个单元格的行号
            int lineNo = Integer.parseInt(firstJavaTableParse.getLineId().split(com.turing.common.util.StringUtils.UNDERLINE)[2]);

            // 不是相邻的2页 或者 第一个单元格的行号不是第一行，不需要合并
            if (page2 != page1 + 1 || lineNo > 1) {
                needMerge = false;
            } else {
                // 通过行号获取行所在的索引
                String lastLineId = lastJavaTableParse.getLineId();
                String firstLineId = firstJavaTableParse.getLineId();
                Integer index1 = lineIndexMap.get(lastLineId);
                Integer index2 = lineIndexMap.get(firstLineId);
                // 如果2行之间的行超过4行，说明不是统一表格，不合并
                if (index1 == null || index2 == null || index2 - index1 > 4) {
                    needMerge = false;
                    // 两行之间差一行（通常是页码那一行）, 合并
                } else if (index2 - index1 == 1) {
                    needMerge = true;
                } else {
                    // 两行之间的行数为2和3时
                    // 判断2个表格的列数是否相同
                    int column1 = Integer.parseInt(lastJavaTableParse.getColumn().split(com.turing.common.util.StringUtils.HORIZONTAL_BAR)[1]) + 1;
                    // 获取当前表格的最后一个单元格，单元格的结束列+1，即时表格的列数
                    JavaTableParses.JavaTableParse curLastJavaParse = curTable.get(curTable.size() - 1);
                    int column2 = Integer.parseInt(curLastJavaParse.getColumn().split(com.turing.common.util.StringUtils.HORIZONTAL_BAR)[1]) + 1;
                    // 如果2个表格的列数相同时，需要合并
                    needMerge = column1 == column2;
                }
            }

            // 合并表格
            if (needMerge) {
                JavaTableParses javaTableParses = mergeTwoTableData(mergedTableList.get(mergedTableList.size() - 1), tableList.get(i));
                mergedTableList.set(mergedTableList.size() - 1, javaTableParses);
            } else {
                mergedTableList.add(tableList.get(i));
            }
        }
        return mergedTableList;
    }

    /**
     * 合并2个表格
     *
     * @param table1 表格1
     * @param table2 表格2
     * @return 返回合并后的表格
     */
    public static JavaTableParses mergeTwoTableData(JavaTableParses table1, JavaTableParses table2) {
        List<JavaTableParses.JavaTableParse> tableParseList = table1.getTable();
        // 通过获取最后一个单元格的行数，来得到本表格的行数
        JavaTableParses.JavaTableParse lastTableParse = tableParseList.get(tableParseList.size() - 1);
        String row = lastTableParse.getRow();
        int rowNum = Integer.parseInt(row.split(com.turing.common.util.StringUtils.HORIZONTAL_BAR)[1]) + 1;

        List<JavaTableParses.JavaTableParse> tableParseList2 = table2.getTable();

        // 对于下面表格的单元格，行数需要加上上一个表格的行数，列数不变
        for (JavaTableParses.JavaTableParse tableParse : tableParseList2) {
            tableParse.setRow(addRowNumToRow(tableParse.getRow(), rowNum));
        }
        tableParseList.addAll(tableParseList2);
        table1.setTable(tableParseList);

        return table1;
    }

    /**
     * 获取新的行字符串
     *
     * @param row    单元格的行字符串
     * @param rowNum 加的行数
     * @return 返回新的行字符串
     */
    public static String addRowNumToRow(String row, Integer rowNum) {
        String[] split = row.split(com.turing.common.util.StringUtils.HORIZONTAL_BAR);
        split[0] = Integer.parseInt(split[0]) + rowNum + "";
        split[1] = Integer.parseInt(split[1]) + rowNum + "";

        return split[0] + com.turing.common.util.StringUtils.HORIZONTAL_BAR + split[1];
    }

    /**
     * 解析表格，将表格中单元格信息进行处理，封装单元格对应的同列和同行的数据
     *
     * @param tableList 表格数据
     * @param ocrJson   ocr结果
     * @return
     */
    public static List<ParseTableVO> parseTable(List<JavaTableParses> tableList, String ocrJson) {
        //执行表格合并
        tableList = mergeTableData(ocrJson, tableList);

        //多张结构化后的表
        List<ParseTableVO> parseTables = new ArrayList<>();

        for (int i = 0; i < tableList.size(); i++) {
            //javaTableParse表示一张表（模型传入）
            JavaTableParses table = tableList.get(i);

            if (table == null) {
                continue;
            }

            //结构化后的表
            List<ParseContentVO> parseTable = new ArrayList<>();

            List<JavaTableParses.JavaTableParse> cells = table.getTable();
            if (CollectionUtils.isEmpty(cells)) {
                continue;
            }

            for (int k = 0; k < cells.size(); k++) {
                //代表每行内容
                JavaTableParses.JavaTableParse cell = cells.get(k);
                //处理单元格，封装同列同行数据
                ParseContentVO attrTable = getAttr(cell, table);
                parseTable.add(attrTable);
            }

            ParseTableVO parseTableVO = new ParseTableVO();
            parseTableVO.setTable(parseTable);
            parseTableVO.setOcr(table.getOcr());
            parseTables.add(parseTableVO);
        }

        return parseTables;
    }


    /**
     * 获取每个单元格向上和向左的数据
     * 规则：
     * 向上找单元格，左的单元格的列范围Z要包括当前单元格的范围D  （z1<=d1,d2<=z2）
     * 向左找单元格，左的单元格的行范围Z要包括当前单元格的范围D  （z1<=d1,d2<=z2）
     *
     * @param cell  单元格数据
     * @param table 表格数据
     * @return
     */
    public static ParseContentVO getAttr(JavaTableParses.JavaTableParse cell, JavaTableParses table) {
        ParseContentVO parseContentVO = new ParseContentVO();
        parseContentVO.setValue(cell.getValue());

        //获得当前值的行与列
        List<Integer> nowColumn = getRowNumOrColumn(cell.getColumn());
        List<Integer> nowRow = getRowNumOrColumn(cell.getRow());

        //该单元格往上的其他单元格数据
        ArrayList<String> topTitle = new ArrayList<>();
        //该单元格往左的其他单元格数据
        ArrayList<String> leftTitle = new ArrayList<>();
        for (int i = 0; i < table.getTable().size(); i++) {
            JavaTableParses.JavaTableParse javaTableParse1 = table.getTable().get(i);
            List<Integer> column = getRowNumOrColumn(javaTableParse1.getColumn());
            List<Integer> row = getRowNumOrColumn(javaTableParse1.getRow());
            //获得行属性(当前行的左边)
            if (row.get(0) <= nowRow.get(0) && row.get(1) >= nowRow.get(1) && column.get(1) < nowColumn.get(0)) {
                //满足中文才取
                if (isTitle(javaTableParse1.getValue())) {
                    leftTitle.add(javaTableParse1.getValue());
                }
            }
            //获得列属性
            if (column.get(0) <= nowColumn.get(0) && column.get(1) >= nowColumn.get(1) && nowRow.get(0) > row.get(1)) {
                //满足中文才取
                if (isTitle(javaTableParse1.getValue())) {
                    topTitle.add(javaTableParse1.getValue());
                }
            }
        }

        parseContentVO.setLeftTitle(leftTitle);
        parseContentVO.setTopTitle(topTitle);
        parseContentVO.setLineId(cell.getLineId());
        parseContentVO.setLinePos(cell.getLinePos());
        parseContentVO.setColumnNum(nowColumn);
        parseContentVO.setRowNum(nowRow);

        return parseContentVO;
    }

    /**
     * 获取当前单元格的起始
     *
     * @param rowOrColumn   行或列的字符串坐标
     *
     * @return  分割后的坐标
     */
    public static List<Integer> getRowNumOrColumn(String rowOrColumn) {
        String[] split = rowOrColumn.split(StringUtils.HORIZONTAL_BAR);
        ArrayList<Integer> result = new ArrayList<>();
        if (split.length == CELL_SPILT_NUM) {
            result.add(Integer.valueOf(split[0]));
            result.add(Integer.valueOf(split[1]));
            return result;
        }
        Log.error("单元格合并信息缺失");
        return result;
    }

    /**
     * bug：表头不是一、开头的，否则不会识别成表头
     */
    public static Boolean isTitle(String text) {
        text = text.replaceAll("\\s*", "");
        boolean flag = false;
        //2个中文：表头； 数值：内容
        String value = "";
        //只匹配中文(空格不算)

        Matcher matcher = MATCH_CHINESE.matcher(text);
        if (matcher.find()) {
            value = matcher.group();
            //空字符串情况除外
            if (org.apache.commons.lang3.StringUtils.isNotBlank(value)) {
                flag = true;
            }
        }

        return flag;
    }

    public static void main(String[] args) {
        String s = "C:\\Users\\98637\\Desktop\\rBIHD2P0XbyACoYXAAEZt_wNR0Q10.json";
        String json = FileUtils.readFile(s);
        List<ParseTableVO> tableVOList = handleTables(JSONArray.parseArray(json));
        System.out.println();
    }
}
